import os
import sys
import cv2
import torch
import numpy as np
import torch.backends.cudnn as cudnn
from pathlib import Path

# 添加车牌识别模块的路径到系统路径
current_dir = os.path.dirname(os.path.abspath(__file__))
vehicle_plate_recognition_dir = os.path.join(current_dir, 'plate_re')
sys.path.append(vehicle_plate_recognition_dir)

# 导入车牌识别所需模块
from plate_re.models.experimental import attempt_load
from plate_re.utils.utils import apply_classifier, non_max_suppression, scale_coords, torch_utils
from plate_re.utils.datasets import letterbox
from plate_re.models.LPRNet import LPRNet, CHARS


class LicensePlateRecognition:
    def __init__(self, 
                 det_weights='plate_re/weights/yolov5_best.pt',
                 rec_weights='plate_re/weights/lprnet_best.pth',
                 img_size=640, 
                 conf_thres=0.4, 
                 iou_thres=0.5, 
                 device=''):
        """
        初始化车牌识别模块
        
        Args:
            det_weights: 检测模型权重路径
            rec_weights: 识别模型权重路径
            img_size: 图像大小
            conf_thres: 置信度阈值
            iou_thres: IOU阈值
            device: 设备选择 (cpu或cuda)
        """
        self.img_size = img_size
        self.conf_thres = conf_thres
        self.iou_thres = iou_thres
        
        # 初始化设备
        self.device = torch_utils.select_device(device)
        self.half = self.device.type != 'cpu'  # 如果是GPU则使用FP16精度
        
        # 加载检测模型
        self.det_model = attempt_load(det_weights, map_location=self.device)
        print("加载检测模型成功!")
        self.img_size = self.check_img_size(img_size, s=self.det_model.stride.max())
        if self.half:
            self.det_model.half()
            
        # 加载识别模型
        self.rec_model = LPRNet(lpr_max_len=8, phase=False, class_num=len(CHARS), dropout_rate=0).to(self.device)
        self.rec_model.load_state_dict(torch.load(rec_weights, map_location=torch.device('cpu')))
        print("加载识别模型成功!")
        self.rec_model.to(self.device).eval()
        
        # 预热模型
        img = torch.zeros((1, 3, self.img_size, self.img_size), device=self.device)
        _ = self.det_model(img.half() if self.half else img) if self.device.type != 'cpu' else None
    
    def check_img_size(self, img_size, s=32):
        # 检查图像大小是否为指定步长的倍数
        new_size = max(int(img_size), s)
        if new_size % s != 0:
            new_size = (new_size // s + 1) * s
        if new_size != img_size:
            print(f'警告: --img-size {img_size} 必须是 {s} 的倍数, 已调整为 {new_size}')
        return new_size
    
    def preprocess_image(self, img_path):
        """处理输入图像"""
        if isinstance(img_path, str):
            # 如果输入是路径，则读取图像
            # 解决中文路径问题
            try:
                # 先尝试直接读取
                img0 = cv2.imread(img_path)
                # 如果读取失败，可能是中文路径问题
                if img0 is None:
                    # 以二进制形式读取文件
                    img_buffer = np.fromfile(img_path, dtype=np.uint8)
                    # 解码为图像矩阵
                    img0 = cv2.imdecode(img_buffer, cv2.IMREAD_COLOR)
                if img0 is None:
                    raise ValueError(f"无法读取图像: {img_path}")
            except Exception as e:
                raise ValueError(f"读取图像时出错: {img_path}, 错误: {str(e)}")
        else:
            # 如果输入是图像数据，直接使用
            img0 = img_path
            
        # 调整图像大小
        img = letterbox(img0, new_shape=self.img_size)[0]
        
        # 将BGR转为RGB，并转为3x416x416格式
        img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB
        img = np.ascontiguousarray(img)
        
        img = torch.from_numpy(img).to(self.device)
        img = img.half() if self.half else img.float()  # uint8 to fp16/32
        img /= 255.0  # 0 - 255 to 0.0 - 1.0
        if img.ndimension() == 3:
            img = img.unsqueeze(0)
            
        return img, img0
    
    def recognize(self, img_path):
        """
        识别图像中的车牌
        
        Args:
            img_path: 图像路径或OpenCV图像对象
            
        Returns:
            results: 包含识别结果的列表，每个元素是一个字典，包含:
                    - box: 检测框坐标 [x1, y1, x2, y2]
                    - confidence: 置信度
                    - plate_number: 识别出的车牌号
        """
        # 预处理图像
        img, im0 = self.preprocess_image(img_path)
        
        # 检测
        pred = self.det_model(img, augment=False)[0]
        # 应用NMS
        pred = non_max_suppression(pred, self.conf_thres, self.iou_thres, classes=None, agnostic=False)
        
        # 应用识别器
        if pred[0] is not None and len(pred[0]):
            pred, plat_num = apply_classifier(pred, self.rec_model, img, [im0])
        else:
            return []
        
        results = []
        for det, lic_plat in zip(pred[0], plat_num):
            *xyxy, conf, cls = det
            
            # 转换为原始图像尺寸
            xmin, ymin, xmax, ymax = scale_coords(img.shape[2:], torch.tensor([xyxy]).view(1, 4), im0.shape).round()[0].cpu().numpy()
            
            # 构建车牌号
            plate_text = ""
            for i in lic_plat:
                plate_text += CHARS[int(i)]
            
            results.append({
                'box': [float(xmin), float(ymin), float(xmax), float(ymax)],
                'confidence': float(conf),
                'plate_number': plate_text
            })
        
        return results


# 使用示例
if __name__ == "__main__":
    # 初始化识别器
    recognizer = LicensePlateRecognition()
    
    # 识别图像
    image_path = "plate_re/demo/images/004504310344827586-86_270-291&474_390&516-390&513_293&516_291&483_387&474-0_0_3_25_25_25_30_30-228-16.jpg"  # 替换为您的图像路径
    results = recognizer.recognize(image_path)
    
    if results:
        # 在图像上绘制结果
        # 使用np.fromfile读取图片以支持中文路径
        try:
            img_buffer = np.fromfile(image_path, dtype=np.uint8)
            img = cv2.imdecode(img_buffer, cv2.IMREAD_COLOR)
        except:
            img = cv2.imread(image_path)
            
        for result in results:
            box = result['box']
            plate_number = result['plate_number']
            confidence = result['confidence']
            
            # 绘制边界框
            cv2.rectangle(img, 
                         (int(box[0]), int(box[1])), 
                         (int(box[2]), int(box[3])), 
                         (0, 255, 0), 2)
            
            # 绘制车牌信息
            text = f"{plate_number} ({confidence:.2f})"
            cv2.putText(img, text, 
                       (int(box[0]), int(box[1]) - 10), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, 
                       (0, 255, 0), 2)
        
        # 保存结果到文件
        output_dir = "results"
        os.makedirs(output_dir, exist_ok=True)
        output_path = os.path.join(output_dir, "result.jpg")
        cv2.imwrite(output_path, img)
        print(f"识别结果已保存至: {output_path}")
        
        # 打印识别结果
        for i, result in enumerate(results):
            print(f"车牌 {i+1}:")
            print(f"  车牌号: {result['plate_number']}")
            print(f"  置信度: {result['confidence']:.2f}")
            print(f"  位置: {result['box']}")
    else:
        print("未检测到车牌") 